# 导入必要的库
import numpy as np
import pandas as pd
from util import createXY
from sklearn.model_selection import train_test_split
import lazypredict
from lazypredict.Supervised import LazyClassifier
import pickle
import logging
import os

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def main():
    logging.info("开始使用 LazyPredict 训练猫狗分类模型")
    
    # 载入和预处理数据
    X, y = createXY(train_folder="D:/haoya/faiss_dog_cat_question-main/data/train", 
                    dest_folder=".", method="flat")
    X = np.array(X).astype('float32')
    y = np.array(y)
    
    logging.info(f"数据形状: X={X.shape}, y={y.shape}")
    
    if X.shape[0] == 0:
        logging.error("未读取到任何图片数据，请检查数据路径和内容。")
        return
    
    # 数据集分割为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
    logging.info("数据集划分完成")
    
    # 使用 LazyClassifier 进行多模型训练和比较
    logging.info("开始使用 LazyPredict 训练和比较多个分类器...")
    
    # 初始化 LazyClassifier
    clf = LazyClassifier(verbose=0, 
                        ignore_warnings=True, 
                        custom_metric=None,
                        predictions=False,
                        random_state=2023)
    
    # 训练并比较所有模型
    models, predictions = clf.fit(X_train, X_test, y_train, y_test)
    
    # 显示所有模型结果
    print("\n" + "="*80)
    print("所有模型性能比较:")
    print("="*80)
    print(models)
    
    # 找到准确率最高的模型
    best_model_name = models['Accuracy'].idxmax()
    best_accuracy = models['Accuracy'].max()
    
    logging.info(f"最佳模型: {best_model_name}, 准确率: {best_accuracy:.4f}")
    
    # 重新训练最佳模型
    logging.info(f"重新训练最佳模型: {best_model_name}")
    
    # 获取最佳模型的类
    from lazypredict.Supervised import _estimators_
    model_class = None
    
    for name, (est, params) in _estimators_.items():
        if name == best_model_name:
            model_class = est
            break
    
    if model_class is None:
        logging.error(f"找不到模型类: {best_model_name}")
        return
    
    # 训练最佳模型
    best_model = model_class()
    best_model.fit(X_train, y_train)
    
    # 评估最终模型
    final_accuracy = best_model.score(X_test, y_test)
    logging.info(f"最终模型准确率: {final_accuracy:.4f}")
    
    # 保存最佳模型
    model_filename = "best_lazypredict_model.pkl"
    with open(model_filename, "wb") as f:
        pickle.dump(best_model, f)
    
    # 保存模型信息
    model_info = {
        'model_name': best_model_name,
        'accuracy': final_accuracy,
        'all_models': models.to_dict()
    }
    
    with open("model_info.pkl", "wb") as f:
        pickle.dump(model_info, f)
    
    # 保存详细的比较结果到CSV
    models.to_csv("lazypredict_model_comparison.csv")
    
    logging.info(f"最佳模型已保存为: {model_filename}")
    logging.info(f"模型比较结果已保存为: lazypredict_model_comparison.csv")
    logging.info(f"模型信息已保存为: model_info.pkl")
    
    # 打印前5个最佳模型
    print("\n" + "="*80)
    print("前5个最佳模型:")
    print("="*80)
    top_models = models.nlargest(5, 'Accuracy')
    for i, (model_name, row) in enumerate(top_models.iterrows(), 1):
        print(f"{i}. {model_name}:")
        print(f"   准确率: {row['Accuracy']:.4f}")
        print(f"   平衡准确率: {row['Balanced Accuracy']:.4f}")
        print(f"   ROC AUC: {row['ROC AUC']:.4f}")
        print(f"   F1 Score: {row['F1 Score']:.4f}")
        print()

if __name__ == '__main__':
    main()